{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import pickle\n",
    "import numpy as np\n",
    "import os\n",
    "import json\n",
    "import random\n",
    "    \n",
    "# load source words\n",
    "source_words_path = os.path.join(os.getcwd(), 'source_words.pkl')\n",
    "with open(source_words_path, 'rb') as f_source_words:\n",
    "    source_words = pickle.load(f_source_words)\n",
    "    \n",
    "# load target words\n",
    "target_words_path = os.path.join(os.getcwd(), 'target_words.pkl')\n",
    "with open(target_words_path, 'rb') as f_target_words:\n",
    "    target_words = pickle.load(f_target_words)\n",
    "    \n",
    "# load label words\n",
    "label_words_path = os.path.join(os.getcwd(), 'label_words.pkl')\n",
    "with open(label_words_path, 'rb') as f_label_words:\n",
    "    label_words = pickle.load(f_label_words)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "945\n",
      "133\n",
      "27\n",
      "1\n",
      "3\n",
      "2\n",
      "0\n",
      "1\n"
     ]
    }
   ],
   "source": [
    "print(len(source_words))\n",
    "print(len(target_words))\n",
    "print(len(label_words))\n",
    "print(source_words['<pad>'])\n",
    "print(source_words['<sep>'])\n",
    "print(source_words['<cls>'])\n",
    "print(source_words['<unk>'])\n",
    "print(target_words['<pad>'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "编码器Encoder的实现\n",
    "'''\n",
    "DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "class Encoder(nn.Module):\n",
    "    def __init__(self, input_dim, hid_dim, n_layers, n_heads, pf_dim, dropout, max_length=100):\n",
    "        super(Encoder, self).__init__()\n",
    "        self.tok_embedding = nn.Embedding(input_dim, hid_dim)\n",
    "        self.pos_embedding = nn.Embedding(max_length, hid_dim)\n",
    "        \n",
    "        # 多层encoder\n",
    "        self.layers = nn.ModuleList([EncoderLayer(hid_dim, n_heads, pf_dim, dropout) for _ in range(n_layers)])\n",
    "        \n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "        self.scale = torch.sqrt(torch.FloatTensor([hid_dim])).to(DEVICE)\n",
    "\n",
    "    def forward(self, src, src_mask):\n",
    "        #src:[batch_size, src_len]\n",
    "        #src_mask:[batch_size, 1, 1, src_len]\n",
    "        \n",
    "        batch_size = src.shape[0]\n",
    "        src_len = src.shape[1]\n",
    "        #位置信息\n",
    "        pos = torch.arange(0, src_len).unsqueeze(0).repeat(batch_size, 1).to(DEVICE)\n",
    "        #token编码+位置编码\n",
    "        src = self.dropout((self.tok_embedding(src) * self.scale) + self.pos_embedding(pos)) # [batch_size, src_len, hid_dim]\n",
    "        \n",
    "        for layer in self.layers:\n",
    "            src = layer(src, src_mask) #[batch_size, src_len, hid_dim]\n",
    "        \n",
    "        return src\n",
    "            \n",
    "class EncoderLayer(nn.Module):\n",
    "    def __init__(self, hid_dim, n_heads, pf_dim, dropout):\n",
    "        super(EncoderLayer, self).__init__()\n",
    "        self.self_attn_layer_norm = nn.LayerNorm(hid_dim)\n",
    "        self.ff_layer_norm = nn.LayerNorm(hid_dim)\n",
    "        self.self_attention = MultiHeadAttentionLayer(hid_dim, n_heads, dropout)\n",
    "        self.positionwise_feedforward = PositionwiseFeedforwardLayer(hid_dim, pf_dim, dropout)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, src, src_mask):\n",
    "        # src:[batch_size, src_len, hid_dim]\n",
    "        # src_mask:[batch_size, 1, 1, src_len]\n",
    "        \n",
    "        # 1.经过多头attetnion后，再经过add+norm\n",
    "        # self-attention\n",
    "        _src = self.self_attention(src, src, src, src_mask)\n",
    "        \n",
    "        src = self.self_attn_layer_norm(src + self.dropout(_src)) # [batch_size, src_len, hid_dim]\n",
    "        \n",
    "        # 2.经过一个前馈网络后，再经过add+norm\n",
    "        _src = self.positionwise_feedforward(src)\n",
    "        \n",
    "        src = self.ff_layer_norm(src + self.dropout(_src)) # [batch_size, src_len, hid_dim]\n",
    "        \n",
    "        return src     \n",
    "\n",
    "class MultiHeadAttentionLayer(nn.Module):\n",
    "    def __init__(self, hid_dim, n_heads, dropout):\n",
    "        super(MultiHeadAttentionLayer, self).__init__()\n",
    "        assert hid_dim % n_heads == 0\n",
    "        \n",
    "        self.hid_dim = hid_dim\n",
    "        self.n_heads = n_heads\n",
    "        self.head_dim = hid_dim // n_heads\n",
    "        \n",
    "        self.fc_q = nn.Linear(hid_dim, hid_dim)\n",
    "        self.fc_k = nn.Linear(hid_dim, hid_dim)\n",
    "        self.fc_v = nn.Linear(hid_dim, hid_dim)\n",
    "        \n",
    "        self.fc_o = nn.Linear(hid_dim, hid_dim)\n",
    "        \n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "        self.scale = torch.sqrt(torch.FloatTensor([self.head_dim])).to(DEVICE)\n",
    "        \n",
    "    def forward(self, query, key, value, mask=None):\n",
    "        batch_size = query.shape[0]\n",
    "        \n",
    "        # query:[batch_size, query_len, hid_dim]\n",
    "        # key:[batch_size, query_len, hid_dim]\n",
    "        # value:[batch_size, query_len, hid_dim]\n",
    "        \n",
    "        Q = self.fc_q(query)\n",
    "        K = self.fc_k(key)\n",
    "        V = self.fc_v(value)\n",
    "        \n",
    "        Q = Q.view(batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3) # [batch_size, query_len, n_heads, head_dim]\n",
    "        K = K.view(batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3)\n",
    "        V = V.view(batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3)\n",
    "        \n",
    "        energy = torch.matmul(Q, K.permute(0, 1, 3, 2)) / self.scale # [batch_size, n_heads, query_len, key_len]\n",
    "        \n",
    "        if mask is not None:\n",
    "            energy = energy.mask_fill(mask == 0, -1e10)\n",
    "        \n",
    "        attention = torch.softmax(energy, dim=-1) # [batch_size, n_heads, query_len, key_len]\n",
    "        \n",
    "        x = torch.matmul(self.dropout(attention), V) # [batch_size, n_heads, query_len, head_dim]\n",
    "        \n",
    "        x = x.permute(0, 2, 1, 3).contiguous() # [batch_size, query_len, n_heads, head_dim]\n",
    "        \n",
    "        x = x.view(batch_size, -1, self.hid_dim) # [batch_size, query_len, hid_dim]\n",
    "        \n",
    "        x = self.fc_o(x) # [batch_size, query_len, hid_dim]\n",
    "        \n",
    "        return x\n",
    "        \n",
    "class PositionwiseFeedforwardLayer(nn.Module):\n",
    "    def __init__(self, hid_dim, pf_dim, dropout):\n",
    "        super(PositionwiseFeedforwardLayer, self).__init__()\n",
    "        self.fc_1 = nn.Linear(hid_dim, pf_dim)\n",
    "        self.fc_2 = nn.Linear(pf_dim, hid_dim)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.gelu = nn.GELU()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        # x:[batch_size, seq_len, hid_dim]\n",
    "        \n",
    "        x = self.dropout(self.gelu(self.fc_1(x))) # [batch_size, seq_len, pf_dim]\n",
    "        x = self.fc_2(x) # [batch_size, seq_len, hid_dim]\n",
    "        \n",
    "        return x\n",
    "\n",
    "class BERT(nn.Module):\n",
    "    def __init__(self, input_dim, hid_dim, n_layers, n_heads, pf_dim, dropout, slot_size, intent_size, src_pad_idx):\n",
    "        super(BERT, self).__init__()\n",
    "        self.src_pad_idx = src_pad_idx\n",
    "        self.encoder = Encoder(input_dim, hid_dim, n_layers, n_heads, pf_dim, dropout)\n",
    "        self.gelu = nn.GELU()\n",
    "        \n",
    "        self.fc = nn.Sequential(nn.Linear(hid_dim, hid_dim), nn.Dropout(dropout), nn.Tanh())\n",
    "        self.intent_out = nn.Linear(hid_dim, intent_size)\n",
    "        self.linear = nn.Linear(hid_dim, hid_dim)\n",
    "       \n",
    "        embed_weight = self.encoder.tok_embedding.weight\n",
    "        self.slot_out = nn.Linear(hid_dim, slot_size, bias=False)\n",
    "        self.slot_out.weight = embed_weight\n",
    "    \n",
    "    def make_src_mask(self, src):\n",
    "        # src: [batch_size, src_len]\n",
    "        src_mask = (src != self.src_pad_idx).unsqueeze(1).unsqueeze(2) # [batch_size, 1, 1, src_len]\n",
    "        \n",
    "    def forward(self, src):\n",
    "        src_mask = self.make_src_mask(src)\n",
    "        encoder_out = self.encoder(src, src_mask) #[batch_size, src_len, hid_dim]\n",
    "        \n",
    "        # 拿到[cls] token进行意图分类\n",
    "        cls_hidden = self.fc(encoder_out[:, 0]) # [batch_size, hid_dim]\n",
    "        intent_output = self.intent_out(cls_hidden) # [batch_size, intent_size]\n",
    "        \n",
    "        # 排除cls进行slot预测\n",
    "        other_hidden = self.gelu(self.linear(encoder_out[:,1:])) # [batch_sze, src_len-1, hid_dim]\n",
    "        slot_output = self.slot_out(other_hidden) # [batch_size, src_len-1, slot_size]\n",
    "        return intent_output, slot_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "n_layers = 6 # transformer-encoder层数\n",
    "n_heads = 12 # 多头self-attention\n",
    "hid_dim =768 \n",
    "dropout = 0.5\n",
    "pf_dim = 768 * 4 \n",
    "\n",
    "model_path = os.path.join(os.getcwd(), \"model.h5\")\n",
    "\n",
    "input_dim = len(source_words) # source 词典大小（即词数量）\n",
    "slot_size = len(target_words) # target 词典大小（即实体类型数量）\n",
    "intent_size = len(label_words) # label 词典大小（即意图类别数量）\n",
    "\n",
    "src_pad_idx = source_words['<pad>']\n",
    "\n",
    "model = BERT(input_dim, hid_dim, n_layers, n_heads, pf_dim, dropout, slot_size, intent_size, src_pad_idx).to(DEVICE)\n",
    "\n",
    "model.load_state_dict(torch.load(model_path))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['<cls>', 'i', 'would', 'like', 'to', 'find', 'a', 'flight', 'from', 'charlotte', 'to', 'las', 'vegas', 'that', 'makes', 'a', 'stop', 'in', 'st.', 'louis', '<sep>']\n",
      "[2, 13, 40, 29, 4, 87, 16, 11, 5, 100, 4, 90, 89, 34, 345, 16, 127, 18, 67, 144, 3]\n",
      "len source :21\n",
      "slot_prediciton:o o o o o o o o o o o o o o o o o o o o\n",
      "intent_prediction:flight\n"
     ]
    }
   ],
   "source": [
    "sentence = \"i would like to find a flight from charlotte to las vegas that makes a stop in st. louis\"\n",
    "sentence2 = \"which airlines have first class flights today\"\n",
    "sentence3 = 'on april first i need a ticket from tacoma to san jose departing before 7 am'\n",
    "sentence4 = 'what are the departure times from detroit to westchester county'\n",
    "sentence5 = 'please find a flight round trip from los angeles to tacoma washington with a stopover in san francisco not exceeding the price of 300 dollars for june tenth 1993'\n",
    "sentence6 = 'show me flight us 1207 from indianapolis to charlotte on monday and flight us 1500 from charlotte to minneapolis on monday and flight twa 639 from minneapolis to indianapolis'\n",
    "sentence7 = 'i need to fly from denver to westchester county on june seventh after 3 pm'\n",
    "sentence8 = 'what meals are served on american flight 811 from tampa to milwaukee'\n",
    "sentence9 = 'meals are served on american flight 665 673 from milwaukee to seattle'\n",
    "\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    tokenized = sentence.split()  # tokenize the sentence\n",
    "    tokenized = ['<cls>'] + tokenized + ['<sep>']\n",
    "    indexed = [source_words[t] for t in tokenized]  # convert to integer sequence\n",
    "    #pad = [1]*(seq_len - len(indexed))\n",
    "    #indexed.extend(pad)\n",
    "    print(tokenized)\n",
    "    print(indexed)\n",
    "    print('len source :{}'.format(len(indexed)))\n",
    "    src_tensor = torch.LongTensor(indexed)  # convert to tensor\n",
    "    src_tensor = src_tensor.unsqueeze(0).to(DEVICE)  # reshape in form of batch,no. of words\n",
    "    \n",
    "    intent_output, slot_output = model(src_tensor)  # prediction\n",
    "    \n",
    "    intent_output = intent_output.squeeze()\n",
    "    intent_output = intent_output.argmax()\n",
    "    intent = intent_output.detach().item()\n",
    "    \n",
    "    pred_token = slot_output.squeeze().argmax(1)\n",
    "    \n",
    "    slot_prediction = [target_words.itos[t.item()] for t in pred_token]\n",
    "    \n",
    "    print('slot_prediciton:{}'.format(' '.join(slot_prediction)))\n",
    "    print('intent_prediction:{}'.format(label_words.itos[intent]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "O O O B-airline_code B-flight_number O B-fromloc.city_name O B-toloc.city_name O B-depart_date.day_name O O O B-flight_number O B-fromloc.city_name O B-toloc.city_name O B-depart_date.day_name O O B-airline_code B-flight_number O B-fromloc.city_name O B-toloc.city_name"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
